"""
国密算法工具类
包含SM2和SM4加解密功能
"""

import binascii
import base64
import random
import string

from gmssl import sm2
from asn1crypto import pem
from sm_crypto import sm2 as sm2_crypto
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
from pyasn1.codec.der import decoder

from jlpay.exception.exceptions import CryptoException, KeyInitializationException

# 常量定义
SM2_PREFIX_BYTE = b'\x04'  # SM2密文前缀
SM4_KEY_LENGTH = 16  # SM4密钥长度，单位字节
DEFAULT_RANDOM_STRING_LENGTH = 32  # 随机字符串默认长度


def parse_public_key(public_key_pem):
    """
    解析公钥PEM格式为16进制字符串
    
    Args:
        public_key_pem: 公钥PEM格式字符串
        
    Returns:
        公钥16进制字符串
    """
    try:
        # 公钥pem格式转16进制
        public_key_bytes = pem.unarmor(public_key_pem.encode())
        try:
            # 尝试使用sm_crypto加载
            public_key_sm2 = sm2_crypto.SM2PublicKey.load(public_key_bytes[2])
            public_key = public_key_sm2.x + public_key_sm2.y
            return public_key
        except Exception:
            # 如果出现异常，使用替代方法加载
            public_key_asn1, _ = decoder.decode(public_key_bytes[2])
            public_key = public_key_asn1[1].asOctets().hex()[2:]
            return public_key
    except Exception as e:
        raise KeyInitializationException("解析公钥格式失败", e)


def parse_private_key(private_key_pem):
    """
    解析私钥PEM格式为16进制字符串
    
    Args:
        private_key_pem: 私钥PEM格式字符串
    
    Returns:
        私钥16进制字符串
    """
    try:
        # 私钥pem格式转16进制
        private_key_bytes = pem.unarmor(private_key_pem.encode())
        private_key_sm2 = sm2_crypto.SM2PrivateKey.load(private_key_bytes[2])
        private_key_asn_hex = private_key_sm2.value
        hex_str_len = private_key_asn_hex[12:14]
        int_val = int(hex_str_len, 16)
        return private_key_asn_hex[14:(14 + int_val * 2)]
    except Exception as e:
        raise KeyInitializationException("解析私钥格式失败", e)


def sm2_sign(message, private_key, public_key):
    """
    使用SM2算法进行签名
    
    Args:
        message: 待签名的消息
        private_key: 私钥16进制字符串
        public_key: 公钥16进制字符串
        
    Returns:
        签名结果的Base64编码字符串
    """
    try:    
        crypt = sm2.CryptSM2(private_key=private_key, public_key=public_key, mode=1, asn1=True)
        message_bytes = _ensure_bytes(message)
        sign_str = crypt.sign_with_sm3(message_bytes)
        return base64.b64encode(bytes.fromhex(sign_str)).decode()
    except Exception as e:
        raise CryptoException("SM2签名失败", e)


def sm2_verify(message, signature, public_key):
    """
    使用SM2算法验证签名
    
    Args:
        message: 原始消息
        signature: Base64编码的签名
        public_key: 公钥16进制字符串
        
    Returns:
        验证结果，True或False
    """
    try:
        crypt = sm2.CryptSM2(public_key=public_key, private_key=None, mode=1, asn1=True)
        sign_bytes = base64.b64decode(signature)
        message_bytes = _ensure_bytes(message)
        return crypt.verify_with_sm3(sign_bytes.hex(), message_bytes)
    except Exception as e:
        raise CryptoException("SM2验签失败", e)


def sm2_encrypt(message, public_key):
    """
    使用SM2算法进行加密
    
    Args:
        message: 待加密的消息
        public_key: 公钥16进制字符串
        
    Returns:
        加密结果的Base64编码字符串
    """
    try:
        # 使用gmssl.sm2中的CryptSM2类进行加密
        sm2_crypt = sm2.CryptSM2(public_key=public_key, private_key=None, mode=1, asn1=True)
        
        # 尝试对消息进行base64解码或转为字节
        message_bytes = _try_base64_decode(message)
            
        # 加密数据
        encrypt_data = sm2_crypt.encrypt(message_bytes)
        # 在加密结果前添加04前缀
        encrypt_data = SM2_PREFIX_BYTE + encrypt_data
        # 返回Base64编码的密文
        return base64.b64encode(encrypt_data).decode()
    except Exception as e:
        raise CryptoException("SM2加密失败", e)


def sm2_decrypt(ciphertext, private_key, public_key):
    """
    使用SM2算法进行解密
    
    Args:
        ciphertext: Base64编码的密文
        private_key: 私钥16进制字符串
        public_key: 公钥16进制字符串

    Returns:
        解密后的明文
    """
    try:
        # 解码Base64密文
        cipher_bytes = base64.b64decode(ciphertext)
        
        # 使用gmssl.sm2中的CryptSM2类进行解密
        sm2_crypt = sm2.CryptSM2(public_key=public_key, private_key=private_key, mode=1, asn1=True)
        
        # 检查密文是否以04开头，如果是则去掉
        if cipher_bytes.startswith(SM2_PREFIX_BYTE):
            cipher_bytes = cipher_bytes[1:]
        # 解密数据
        decrypt_data = sm2_crypt.decrypt(cipher_bytes)
        
        # 将解密后的数据进行base64编码后返回
        return base64.b64encode(decrypt_data).decode()
    except Exception as e:
        raise CryptoException("SM2解密失败", e)


def sm4_ecb_encrypt(text, key):
    """
    SM4/ECB/PKCS5Padding加密算法
    
    Args:
        text: 待加密明文
        key: SM4对称加密密钥, base64str
        
    Returns:
        加密后的密文，base64str
    """
    try:
        data = _ensure_bytes(text)
        key_bytes = base64.b64decode(key)
        
        # 验证密钥长度
        if len(key_bytes) != SM4_KEY_LENGTH:
            raise ValueError(f"SM4 Key长度必须为{SM4_KEY_LENGTH}字节")
            
        crypt_sm4 = CryptSM4()
        crypt_sm4.set_key(key_bytes, SM4_ENCRYPT)
        encrypt = crypt_sm4.crypt_ecb(data)
        return base64.b64encode(encrypt).decode("utf-8")
    except Exception as e:
        raise CryptoException("SM4加密失败", e)


def sm4_ecb_decrypt(content, key):
    """
    SM4/ECB/PKCS5Padding解密算法    
    Args:
        content: 密文
        key: SM4对称加密密钥
        
    Returns:
        解密后的明文
    """
    try:
        data = base64.b64decode(_ensure_bytes(content))
        key_bytes = base64.b64decode(key)
        crypt_sm4 = CryptSM4()
        crypt_sm4.set_key(key_bytes, SM4_DECRYPT)
        decrypt = crypt_sm4.crypt_ecb(data)
        return decrypt.decode("utf-8")
    except Exception as e:
        raise CryptoException("SM4解密失败", e)


def generate_random_string(length=DEFAULT_RANDOM_STRING_LENGTH):
    """
    生成指定长度的随机字符串
    
    Args:
        length: 随机字符串长度，默认32
        
    Returns:
        随机字符串
    """
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))


def generate_sm4_key():
    """
    生成SM4密钥
    
    Returns:
        Base64编码的SM4密钥
    """
    try:
        # 生成16字节随机数作为SM4密钥
        key_bytes = bytes(random.getrandbits(8) for _ in range(SM4_KEY_LENGTH))
        return base64.b64encode(key_bytes).decode('utf-8')
    except Exception as e:
        raise CryptoException("生成SM4密钥失败", e)


def parse_sm2_key(private_key_pem=None, public_key_pem=None):
    """
    解析SM2密钥PEM格式为16进制字符串
    
    Args:
        private_key_pem: 私钥PEM格式字符串，可选
        public_key_pem: 公钥PEM格式字符串，可选
        
    Returns:
        如果提供私钥，返回(私钥16进制字符串, None)
        如果提供公钥，返回(None, 公钥16进制字符串)
        如果同时提供私钥和公钥，返回(私钥16进制字符串, 公钥16进制字符串)
    
    Raises:
        KeyInitializationException: 解析密钥格式失败时抛出
    """
    pri_key_hex = None
    pub_key_hex = None
    
    if private_key_pem is not None:
        pri_key_hex = parse_private_key(private_key_pem)
    
    if public_key_pem is not None:
        pub_key_hex = parse_public_key(public_key_pem)
    
    return pri_key_hex, pub_key_hex

def get_digest_by_type(digest_type, data):
    """
    根据摘要算法类型生成摘要
    
    Args:
        digest_type: 摘要算法类型，支持 'SM3' 和 'SHA256'
        data: 需要计算摘要的数据，可以是字符串或字节
        
    Returns:
        Base64编码的摘要值
        
    Raises:
        CryptoException: 计算摘要失败时抛出
    """
    try:
        # 确保数据为字节类型
        data_bytes = _ensure_bytes(data)
        
        # 根据算法类型计算摘要
        if digest_type == 'SM3':
            from gmssl import sm3
            # 将bytes转换为bytearray进行处理
            data_list = [b for b in data_bytes]
            # sm3_hash返回的是十六进制字符串，需要转换为字节
            hex_str = sm3.sm3_hash(data_list)
            message = bytes.fromhex(hex_str)
        elif digest_type == 'SHA256':
            import hashlib
            message = hashlib.sha256(data_bytes).digest()
        else:
            raise CryptoException(f"不支持的摘要算法类型: {digest_type}")
            
        # Base64编码
        return base64.b64encode(message).decode('utf-8')
    except Exception as e:
        raise CryptoException(f"计算摘要失败: {str(e)}", e)


# 私有辅助函数
def _ensure_bytes(data):
    """
    确保数据为字节类型
    
    Args:
        data: 字符串或字节
        
    Returns:
        字节类型数据
    """
    if isinstance(data, str):
        return data.encode()
    return data


def _try_base64_decode(message):
    """
    尝试对消息进行base64解码，如果失败则返回原始消息
    
    Args:
        message: 原始消息
        
    Returns:
        解码后的消息或原始消息
    """
    if isinstance(message, str):
        try:
            return base64.b64decode(message)
        except:
            return message.encode()
    return message 